from langchain.prompts import PromptTemplate
from langchain.prompts import FewShotPromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector
from langchain.chains import LLMChain
import os

from langchain_community.chat_models import ChatTongyi

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
# 定义模型
llm = ChatTongyi()

# 创建一些例子
examples = [
    {"input": "happy", "output": "sad"},
    {"input": "tall", "output": "short"},
    {"input": "energetic", "output": "lethargic"},
    {"input": "sunny", "output": "gloomy"},
    {"input": "windy", "output": "calm"},
]

example_prompt = PromptTemplate(
    input_variables=["input", "output"],
    template="Input: {input}\nOutput: {output}",
)
# 调用长度选择器 类型 LengthBasedExampleSelector
example_selector = LengthBasedExampleSelector(
    # 可供选择的示例。
    examples=examples,
    # 用于格式化示例的PromptTemplate。
    example_prompt=example_prompt,
    # 格式化示例的最大长度。
    # 长度由下面的get_text_Length函数测量。
    max_length=25,
    # 用于获取字符串长度的函数
    # 以确定要包括哪些示例。它被注释掉是因为
    # 如果未指定，则将其作为默认值提供。
    # get_text_length:可调用[[str]，int]=lambda x:len（re.split（“\n|”，x））
)
dynamic_prompt = FewShotPromptTemplate(
    # 提供了ExampleSelector而不是示例。
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="给出每个输入的反义词",
    suffix="Input: {adjective}\nOutput:",
    input_variables=["adjective"],
)

print(dynamic_prompt.format(adjective="happy"))
chain = LLMChain(llm=llm, prompt=dynamic_prompt)
res = chain.invoke("fun")
print(f"the result is {res}")
